import time
import requests
import threading
from enum import Enum
from typing import Union
from motor.utils.singleton import ThreadSafeSingleton
from motor.config.coordinator import CoordinatorConfig
from motor.resources.instance import Instance, PDRole, Endpoint
from motor.coordinator.core.instance_manager import InstanceManager
from motor.utils.logger import get_logger

logger = get_logger(__name__)

class MetricType(Enum):
    GAUGE = "gauge"
    COUNTER = "counter"
    HISTOGRAM = "histogram"
    SUMMARY = "summary"

    @classmethod
    def from_string(cls, type_string):
        try:
            return cls[type_string.upper()]
        except KeyError:
            raise ValueError(f"[Metrics] Invalid type string: {type_string}")

    def __str__(self):
        return self.value

class SingleMetric():
    def __init__(self, single_metric=None):
        if single_metric:
            self.name: str = single_metric.name
            self.help: str = single_metric.help
            self.type: MetricType = single_metric.type
            self.label: list[str] = single_metric.label
            self.value: list[float] = [0.0] * len(single_metric.label)
        else:
            self.name: str = ""
            self.help: str = ""
            self.type: str = ""
            self.label: list[str] = []
            self.value: list[float] = []


class MetricsCollector(ThreadSafeSingleton):
    def __init__(self):
        # If the metrics collector is already initialized, return.
        if hasattr(self, '_initialized'):
            return

        # Initial metrics state
        self._inactive_instance_metrics_aggregate: list[SingleMetric] = []
        self._instance_metrics_cached: dict[int, dict[str, list[SingleMetric]]] = {}
        self._last_metrics: str = None
        self._last_instance_metrics: dict[int, list[SingleMetric]] = None

        self._reuse_time: int = CoordinatorConfig().prometheus_metrics_config.reuse_time
        self._lock = threading.Lock()
        self.stop_event = threading.Event()
        self._metrics_update_thread = threading.Thread(
            target=self.update_metrics,
            daemon=True,
            name="MetricsUpdate"
        )
        self._metrics_update_thread.start()

        self._initialized = True
        logger.info("MetricsCollector started.")

    def update_metrics(self):
        metrics, instance_metrics = self._get_and_aggregate_metrics()

        with self._lock:
            if metrics and instance_metrics:
                self._last_metrics = metrics
                self._last_instance_metrics = instance_metrics

    def _update_metrics_thread(self):
        while not self.stop_event.is_set():
            self.update_metrics()
            time.sleep(self._reuse_time)

    def stop(self) -> None:
        self.stop_event.set()
        self._metrics_update_thread.join()
        logger.info("MetricsCollector stopped.")

    def prometheus_instance_metrics_handler(self):
        with self._lock:
            instance_metrics = self._last_metrics
        return instance_metrics

    def prometheus_metrics_handler(self):
        with self._lock:
            metrics = self._last_metrics
        return metrics

    def _get_server_metrics_single(self, ip: str, port: str) -> str:
        """
        Get metrics of single engine server.

        :param ip: engine server ip
        :param port: engine server port
        :returns: engine server metrics. If request failed, return "".
        """

        url = f"http://{ip}:{port}/metrics"
        try:
            response = requests.get(url)

            if response.status_code == 200:
                data = response.json()
                return data
            else:
                logger.warning(f"[Metrics] request metrics failed: code = {response.status_code}")

        except requests.exceptions.RequestException as e:
            logger.warning(f"[Metrics] request metrics failed: {e}")

        return ""

    def _get_server_metrics(self, available_instances: dict[int, Instance]) -> dict[int, dict]:
        """
        Get instances/endpoints info and get all endpoints metrics text.

        :param available_instances: alive instances
        :returns:
            for example:
            {
                instance_id0: {
                    "endpoints": {
                        endpoint_id0: {
                            "metrics_str": "xxx"
                        },
                        endpoint_id1: ...
                    }
                },
                instance_id1: ...
            }
        """

        collects = {}
        for ins_info in available_instances.values():
            ins_id = ins_info.id
            ens_metrics_ok = True
            collect = {
                "endpoints": {}
            }
            for ens_info in ins_info.endpoints.values():
                for en_info in ens_info.values():
                    metrics_str = self._get_server_metrics_single(en_info.ip, en_info.port)
                    if not metrics_str:
                        ens_metrics_ok = False
                        break
                    collect["endpoints"][en_info.id] = {
                        "metrics_str": metrics_str
                    }
                if not ens_metrics_ok:
                    break
            if ens_metrics_ok:
                collects[ins_id] = collect

        return collects

    def _parse_metric_help(self, single_metric: SingleMetric, line: str) -> bool:
        """
        Parse help line.

        :param single_metric:
        :param line: format: "# HELP <name> <help description>"
        :returns: if success, update single_metric and return True, else return False
        """

        parts = line.split()
        if len(parts) >= 4 and parts[0] == "#" and parts[1] == "HELP":
            single_metric.name = parts[2]
            single_metric.help = " ".join(parts[3:])
            return True
        else:
            logger.error("[Metrics] Parse metric help failed.")
            return False

    def _parse_metric_type(self, single_metric: SingleMetric, line: str) -> bool:
        """
        Parse type line.

        :param single_metric:
        :param line: format: "# TYPE <name> <gauge|counter|histogram|summary>"
        :returns: if success, update single_metric and return True, else return False
        """

        parts = line.split()
        if len(parts) == 4 and parts[0] == "#" and parts[1] == "TYPE" and parts[3] in ["counter", "gauge", "histogram", "summary"]:
            try:
                single_metric.type = MetricType.from_string(parts[3])
            except ValueError:
                logger.error("[Metrics] Illegal metric type: %s", parts[3])
                return False
            return True
        else:
            logger.error("[Metrics] Parse metric type failed.")
            return False

    def _parse_metric_body_block(self, single_metric: SingleMetric, line: str) -> bool:
        """
        Parse help line.

        :param single_metric:
        :param line: format: "<label> <value>"
        :returns: if success, update single_metric and return True, else return False
        """

        parts = line.split()
        if len(parts) == 2:
            single_metric.label.append(parts[0])
            try:
                value = float(parts[1])
                if value < 0:
                    logger.error("[Metrics] Illegal metric value: %s", parts[1])
                    return False
                single_metric.value.append(value)
            except ValueError:
                logger.error("[Metrics] Illegal metric value: %s", parts[1])
                return False
            return True
        else:
            logger.error("[Metrics] Parse metric body failed.")
            return False

    def _parse_metric_text(self, metrics_str) -> list[SingleMetric]:
        """
        Parse metrics from text to format data.

        :param metrics_str:
            metrics_str format:
                # HELP <name> <help description>
                # TYPE <name> <gauge|counter|histogram|summary>
                <label0> <value0>
                <label1> <value1>
                ...
                <labeln> <valuen>
                # HELP ...
                # TYPE ...
                ...
        :returns: If success, return a list of SingleMetric, else return []
            SingleMetric format:
                name = <name>
                help = <help description>
                type = <gauge|counter|histogram|summary>
                label = [label0, label1, ..., labeln]
                value = [value0, value1, ..., valuen]
        """

        metric_array = []
        lines = metrics_str.strip().split("\n")
        i = 0
        while i < len(lines):
            single_metric = SingleMetric()
            if i < len(lines) and not self._parse_metric_help(single_metric, lines[i]):
                return []
            i += 1
            if i < len(lines) and not self._parse_metric_type(single_metric, lines[i]):
                return []
            i += 1
            while i < len(lines) and lines[i][0] != "#":
                if not self._parse_metric_body_block(single_metric, lines[i]):
                    return []
                i += 1
            metric_array.append(single_metric)
        return metric_array

    def _parse_metrics(
        self,
        collects: dict[int, dict[str, dict[int, dict[str, str]]]]
    ) -> bool:
        """
        Parse metrics text to format data for all instances/endpoints.

        :param collects:
            collects before call:
            {
                instance_id0: {
                    "endpoints": {
                        endpoint_id0: {
                            "metrics_str": "xxx"
                        },
                        endpoint_id1: ...
                    }
                },
                instance_id1: ...
            }
            collects after call:
            {
                instance_id0: {
                    "endpoints": {
                        endpoint_id0: {
                            "metrics": metrics_value # the type of metrics_value: list[SingleMetric]
                        },
                        endpoint_id1: ...
                    }
                },
                instance_id1: ...
            }
        :returns: if success, replace "metrics_str" by "metrics" and return True, else return False
        """

        if not isinstance(collects, dict):
            logger.error("[Metrics] Invalid pods metric JSON file.")
            return False

        if not collects:
            return True

        metric_count = 0
        for instance_id in collects.keys():
            if not isinstance(collects[instance_id], dict) or not collects[instance_id] or "endpoints" not in collects[instance_id]:
                logger.error("[Metrics] Invalid pods metric JSON file.")
                return False

            pods = collects[instance_id]["endpoints"]
            for pod_info in pods.values():
                if "metrics_str" not in pod_info:
                    logger.error("[Metrics] Invalid 'metrics_str' in pod metrics JSON file.")
                    return False
                parsed_metric = self._parse_metric_text(pod_info["metrics_str"])
                if metric_count == 0:
                    metric_count = len(parsed_metric)
                elif metric_count != len(parsed_metric):
                    parsed_metric = {}
                if not parsed_metric:
                    logger.error("[Metrics] Parse metric text failed.")
                    return False
                pod_info["metrics"] = parsed_metric
        return True

    def _aggregate_metric_by_sum(
        self,
        endpoints: dict[int, dict[str, list[SingleMetric]]],
        single_metric: SingleMetric,
        index: int
    ) -> None:
        """
        Aggregate the index-th metric of all pods in single instance.

        :param endpoints:
            endpoints format:
            {
                endpoint_id0: {
                    "metrics": metrics_value # the type of metrics_value: list[SingleMetric]
                },
                endpoint_id1: ...
            }
        :param single_metric: aggregate result save to single_metric
        :param index: the position in list[SingleMetric]
        :returns:
        """
        first_pod = next(iter(endpoints.values()))
        value_num = len(first_pod["metrics"][index].value)
        single_metric.value = [0.0] * value_num
        for i in range(value_num):
            for pod_info in endpoints.values():
                single_metric.value[i] += pod_info["metrics"][index].value[i]

    def _aggregate_metrics_by_instance(
        self,
        collects: dict[int, dict[str, dict[int, dict[str, list[SingleMetric]]]]]
    ) -> bool:
        """
        For each instance, aggreagte metrics of all endpoints.

        :param collects:
            collects before call:
            {
                instance_id0: {
                    "endpoints": {
                        endpoint_id0: {
                            "metrics": metrics_value # the type of metrics_value: list[SingleMetric]
                        },
                        endpoint_id1: ...
                    }
                },
                instance_id1: ...
            }
            collects after call:
            {
                instance_id0: {
                    "metrics": metrics_value # the type of metrics_value: list[SingleMetric]
                },
                instance_id1: ...
            }
        :returns:
        """

        for instance_id in collects.keys():
            endpoints = collects[instance_id]["endpoints"]
            if not endpoints:
                continue

            first_pod = next(iter(endpoints.values()))
            aggregate = []
            metric_count = len(first_pod["metrics"])
            for i in range(metric_count):
                single_metric = SingleMetric(first_pod["metrics"][i])
                self._aggregate_metric_by_sum(endpoints, single_metric, i)
                aggregate.append(single_metric)
            collects[instance_id]["metrics"] = aggregate
            del collects[instance_id]["endpoints"]

            if not self._check_and_update_metrics_cached(instance_id, collects[instance_id]["metrics"]):
                logger.error("[Metrics] Update metrics state failed.")
                return False

        return True

    def check_metrics_format(self, metrics: list[SingleMetric]) -> bool:
        """
        Check metrics format.

        :param metrics:
        :returns:
        """
        if not isinstance(metrics, list) or not metrics:
            return False
        if not self._instance_metrics_cached:
            return True
        base_metrics = next(iter(self._instance_metrics_cached.values()))["metrics"]
        if len(base_metrics) != len(metrics):
            return False

        for i in range(len(base_metrics)):
            if base_metrics[i].name != metrics[i].name or \
                base_metrics[i].help != metrics[i].help or \
                base_metrics[i].type != metrics[i].type or \
                base_metrics[i].label != metrics[i].label:
                return False

        return True

    def _check_and_update_metrics_cached(self, instance_id: int, instance_metrics: list[SingleMetric]) -> bool:
        """
        Check metrics format and add/update cache.

        :param instance_id: instance id
        :param instance_metrics: instance metrics
        :returns:
        """

        if not self.check_metrics_format(instance_metrics):
            return False
        self._instance_metrics_cached[instance_id] = {
            "metrics": instance_metrics
        }
        return True

    def _aggregate_instance_metric_by_sum(
        self,
        collects: dict[int, dict[str, list[SingleMetric]]],
        single_metric: SingleMetric,
        index: int
    ) -> None:
        """
        Aggregate the index-th metric of all instance.

        :param collects:
        :param single_metric: aggregate result save to single_metric
        :param index: the position in list[SingleMetric]
        :returns:
        """

        if not self._instance_metrics_cached:
            return

        first_instance = next(iter(self._instance_metrics_cached.values()))
        value_num = len(first_instance["metrics"][index].value)
        for ins_id in self._instance_metrics_cached:
            # gauge type only aggregate active instances
            if single_metric.type == MetricType.GAUGE and ins_id not in collects:
                continue

            for i in range(value_num):
                single_metric.value[i] += self._instance_metrics_cached[ins_id]["metrics"][index].value[i]
                if self._inactive_instance_metrics_aggregate:
                    single_metric.value[i] += self._inactive_instance_metrics_aggregate[index].value[i]

    def _aggregate_metrics_all_instance(self, collects: dict[int, dict[str, list[SingleMetric]]]) -> list[SingleMetric]:
        """
        Aggreagte metrics of all instances.

        :param collects:
        :returns:
        """

        if not self._instance_metrics_cached:
            return []

        first_instance = next(iter(self._instance_metrics_cached.values()))
        aggregate = []
        metric_count = len(first_instance["metrics"])
        for i in range(metric_count):
            single_metric = SingleMetric(first_instance["metrics"][i])
            self._aggregate_instance_metric_by_sum(collects, single_metric, i)
            aggregate.append(single_metric)

        return aggregate

    def _get_value_str(self, value: float) -> str:
        """
        Transform float to str.

        :returns:
        """

        if value == float("nan"):
            return "Nan"
        elif value == float("inf"):
            return "+Inf"
        elif value == float("-inf"):
            return "-Inf"
        return str(value)

    def _get_serialize_metrics(self, aggregate: list[SingleMetric]) -> str:
        """
        Metrics serialize.

        :param aggregate:
        :returns:
        """

        lines = []
        for item in aggregate:
            lines.append("# HELP {} {}".format(item.name, item.help))
            lines.append("# TYPE {} {}".format(item.name, item.type))
            for i in range(len(item.label)):
                lines.append("{} {}".format(item.label[i], self._get_value_str(item.value[i])))
        return "\n".join(lines)

    def _get_serialize_instance_metrics(
        self,
        collects: dict[int, dict[str, list[SingleMetric]]]
    ) -> dict[int, list[SingleMetric]]:
        """
        Instance metrics serialize.

        :param collects:
        :returns:
        """

        instance_metrics = {}
        for ins_id in collects.keys():
            instance_metrics[ins_id] = self._instance_metrics_cached[ins_id]["metrics"]

        return instance_metrics

    def _clear_inactive_metrics(self, unavailable_pool: dict[int, Instance]) -> None:
        clear_ins_list = []
        for ins_id in unavailable_pool.keys():
            if ins_id in self._instance_metrics_cached:
                clear_ins_list.append(ins_id)
        
        for ins_id in clear_ins_list:
            metrics = self._instance_metrics_cached[ins_id]["metrics"]

            # initialize self._inactive_instance_metrics_aggregate
            if not self._inactive_instance_metrics_aggregate:
                for index in range(len(metrics)):
                    aggregate_metric = SingleMetric(metrics[index])
                    self._inactive_instance_metrics_aggregate.append(aggregate_metric)

            # aggregate metrics[index] to self._inactive_instance_metrics_aggregate
            for index in range(len(metrics)):
                if metrics[index].type == MetricType.GAUGE:
                    continue
                for i in range(len(metrics[index].value)):
                    self._inactive_instance_metrics_aggregate[index].value[i] += metrics[index].value[i]

            # remove ins_id from cache
            del self._instance_metrics_cached[ins_id]


    def _get_and_aggregate_metrics(self) -> tuple[str, dict[str, list[SingleMetric]]]:
        """
        Get and Aggregate metrics.

        :returns:
        """

        # Step 0: get instances/endpoints info
        available_instances, unavailable_instances = InstanceManager().get_all_instances()
        self._clear_inactive_metrics(unavailable_instances)

        # Step 1: get instances/endpoints info and get all endpoints metrics text.
        collects = self._get_server_metrics(available_instances)

        # Step 2: parse metrics text to format data for all instances/endpoints.
        if not self._parse_metrics(collects):
            logger.error("[Metrics] Parse vllm server metrics failed.")
            return "", {}

        # Step 3: for each instance, aggreagte metrics of all endpoints.
        self._aggregate_metrics_by_instance(collects)

        # Step 4: aggreagte metrics of all instances.
        aggregate = self._aggregate_metrics_all_instance(collects)

        # Step 5: serialize and return to handler.
        return self._get_serialize_metrics(aggregate), self._get_serialize_instance_metrics(collects)

